#include <iostream>
#include <vector>
#include <unordered_map>
#include <memory>
#include <string>

// Forward declarations
struct Expression;
struct Statement;

// Represents a unique expression using a string identifier
struct Expression {
    std::string id;        // Identifier for the expression (e.g., "a + b")
    std::vector<std::string> operands; // Operands of the expression
    Expression(const std::string& id, const std::vector<std::string>& ops)
        : id(id), operands(ops) {}
};

// Represents an assignment statement
struct Statement {
    std::string result;     // The variable that stores the result of the expression
    std::shared_ptr<Expression> expression; // The expression being assigned
    Statement(const std::string& res, std::shared_ptr<Expression> expr)
        : result(res), expression(expr) {}
};

// Represents a function containing multiple statements
struct Function {
    std::vector<Statement> statements; // List of statements in the function
    std::vector<std::string> returns;  // Variables to return from the function

    void addStatement(const Statement& stmt) {
        statements.push_back(stmt);
    }

    void addReturn(const std::string& ret) {
        returns.push_back(ret);
    }
};

// Represents the entire program
struct Program {
    std::vector<Function> functions; // List of functions in the program

    void addFunction(const Function& func) {
        functions.push_back(func);
    }
};

// CSE Pass class to perform common subexpression elimination
class CSEPass {
public:
    // This function processes the program and eliminates common expressions
    void run(Program& program) {
        for (auto& func : program.functions) {
            eliminateCommonSubexpressions(func);
        }
    }

private:
    // Map to store unique expressions and their assigned variables
    std::unordered_map<std::string, std::string> expressionMap;

    // Eliminate common subexpressions in a function
    void eliminateCommonSubexpressions(Function& func) {
        for (auto& stmt : func.statements) {
            // Check if the expression is already in the map
            std::string exprId = stmt.expression->id;
            if (expressionMap.find(exprId) != expressionMap.end()) {
                // Replace the current statement's result with the existing one
                std::cout << "Replacing " << stmt.result << " with " << expressionMap[exprId] << ", exprId = " << exprId << std::endl;
                stmt.result = expressionMap[exprId];
            } else {
                // Add the expression to the map
                expressionMap[exprId] = stmt.result;
            }
        }

        // For the returns, replace any with their corresponding expressions
        for (auto& ret : func.returns) {
            if (expressionMap.find(ret) != expressionMap.end()) {
                ret = expressionMap[ret];
            }
        }
    }
};

// Example usage
int main() {
    // Create a program with a function
    Program program;

    // Create a function
    Function func;

    // Create statements
    auto expr1 = std::make_shared<Expression>("a + b", std::vector<std::string>{"a", "b"});
    Statement stmt1("c0", expr1);
    
    auto expr2 = std::make_shared<Expression>("a + b", std::vector<std::string>{"a", "b"});
    Statement stmt2("c1", expr2);
    
    auto expr3 = std::make_shared<Expression>("c0 + c", std::vector<std::string>{"c0", "c"});
    Statement stmt3("c2", expr3);
    
    auto expr4 = std::make_shared<Expression>("c2", std::vector<std::string>{"c2"});
    Statement stmt4("c3", expr4);
    
    // Add statements to the function
    func.addStatement(stmt1);
    func.addStatement(stmt2);
    func.addStatement(stmt3);
    func.addStatement(stmt4);
    
    // Add return values
    func.addReturn("c0");
    func.addReturn("c1");
    func.addReturn("c2");
    func.addReturn("c3");
    
    // Add the function to the program
    program.addFunction(func);

    // Run the CSE pass
    CSEPass cse;
    cse.run(program);

    // Print the results
    for (const auto& func : program.functions) {
        for (const auto& stmt : func.statements) {
            std::cout << stmt.result << " = " << stmt.expression->id << std::endl;
        }
        std::cout << "return ";
        for (const auto& ret : func.returns) {
            std::cout << ret << " ";
        }
        std::cout << std::endl;
    }

    return 0;
}
